import collections
import os

import cv2
# import mayavi.mlab as mlab
import numpy as np
import torch

filePath = '/media/st/data_fast/tj/0119/training/image_2'
image_names = sorted(os.listdir(filePath))
lines = []
for image_name in image_names:
    lines.append(filePath+"/"+image_name)

image_num = 0
image_class_ =1
label = []
img = np.zeros(1)

def OnMouse(event,x,y,flags,param):
    global image_class_
    if image_class_ == -1:
        return
    if event==cv2.EVENT_LBUTTONDOWN:#左键单击,这里可以设置为其他的鼠标动作

        label.append([image_class_,x, y ])
        if(image_class_ == 0):
            color=(0,0,255)
            cv2.circle(img, (x, y), 2, color, 2)
            image_class_ = 1
        elif(image_class_ == 1):
            color = (0,0,0)
            cv2.circle(img, (x, y), 2, color, 2)
            image_class_ =  -1
        cv2.imshow('image',img)


cv2.namedWindow('image')
cv2.setMouseCallback('image',OnMouse)

for i  in range(7,len(lines)):
    annotation_line = lines[i]
    line = annotation_line.split()

    img = cv2.imread(line[0])

    cv2.imshow('image',img)
    key = cv2.waitKey(0)
    if key == "q":
        break

    if (len(label) !=  0):
        file = open('/media/st/data_fast/tj/0119/training/image_label.txt', 'a+') 
        # file.write(line[0]+" "+str(label[0][1])+","+str(label[0][2])+","+str(label[0][0])+" "+\
        #     str(label[1][1])+","+str(label[1][2])+","+str(label[1][0])+"\n")
        file.write(line[0]+" "+str(label[0][1])+","+str(label[0][2])+","+str(label[0][0])+" "+"\n")
        file.close()

    image_class_ = 1
    label = []

    print("image num: ",i)



